import torch
import torchvision 
from tqdm import tqdm
#直接引入matplotlib包是不能引入pyplot模块的
#import matplotlib 会出错
import matplotlib.pyplot as plt
import matplotlib as mpl
from PIL import Image
from first import Net


def predict(model, src):
	# 将输入的图像从array格式转为image
    img = Image.open(src)
    # 自己定义的pytorch transform方法
    transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(), torchvision.transforms.Normalize(mean = [0.5],std = [0.5])])
    img = transform(img)
    # .view()用来增加一个维度
    # 我的图像的shape为(1, 64, 64)
    # channel为1，H为64， W为64
    # 因为训练的时候输入的照片的维度为(batch_size, channel, H, W) 所以需要再新增一个维度
    # 增加的维度是batch size，这里输入的一张图片，所以为1
    img = img.view(1, 1, 28, 28)
    output = model(img)
    _, prediction = torch.max(output, 1)
    #将预测结果从tensor转为array，并抽取结果
    prediction = prediction.numpy()[0]
    return prediction


if __name__ == '__main__':
    model = torch.load('./model.pth', map_location=torch.device('cpu'))
    torch.no_grad()
    f = './DataImages-Test/9999-label-6.png'
    res = predict(model, f)
    print(f)
    print(res)

